{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PydanticOutputParser\n",
    "\n",
    "- Author: [Jaeho Kim](https://github.com/Jae-hoya)\n",
    "- Peer Review: [Donghak Lee](https://github.com/stsr1284), [frimer](https://github.com/brian604)\n",
    "- Proofread : [BokyungisaGod](https://github.com/BokyungisaGod)\n",
    "- This is a part of [LangChain Open Tutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/03-OutputParser/01-PydanticOutputParser.ipynb)",
    "[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/03-OutputParser/01-PydanticOutputParser.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "\n",
    "This tutorial covers how to perform ```PydanticOutputParser``` using ```pydantic```.\n",
    "\n",
    "The ```PydanticOutputParser``` is a class that helps transform the output of a language model into **structured information** . This class can **provide the information you need in a clear and organized form** instead of a simple text response.\n",
    "\n",
    "\n",
    "By utilizing this class, you can transform the output of your language model to fit a specific data model, making it easier to process and utilize the information\n",
    "\n",
    "### Main Method\n",
    "\n",
    "A ```PydanticOutputParser``` primarily requires the implementation of **two core methods**.\n",
    "\n",
    "\n",
    "1. **```get_format_instructions()```**: Provide instructions that define the format of the information that the language model should output. \n",
    "For example, you can return instructions as a string that describes the fields of data that the language model should output and how they should be formatted. \n",
    "These instructions are very important for the language model to structure the output and transform it to fit your specific data model.\n",
    "\n",
    "2. **```parse()```**: Takes the output of the language model (assumed to be a string) and analyzes and transforms it into a specific structure. \n",
    "Use a tool like Pydantic to validate the input string against a predefined schema and transform it into a data structure that follows that schema.\n",
    "\n",
    "\n",
    "### Table of Contents\n",
    "- [Overview](#overview)\n",
    "- [Environment Setup](#environment-setup)\n",
    "- [A Chain with the Basic StrOutputParser](#a-chain-with-the-basic-stroutputparser)\n",
    "- [Using the PydanticOutputParser](#using-the-pydanticoutputparser)\n",
    "\n",
    "### References\n",
    "\n",
    "- [Pydantic Official Document](https://docs.pydantic.dev/latest/)\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setups, useful functions, and utilities for tutorials. \n",
    "- You can check out the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain-opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langsmith\",\n",
    "        \"langchain\",\n",
    "        \"langchain_core\",\n",
    "        \"langchain_openai\",\n",
    "        \"pydantic\",\n",
    "        \"itertools\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment variables have been set successfully.\n"
     ]
    }
   ],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"OPENAI_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"01-PydanticOuputParser\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Environment variables have been set successfully.\n",
    "You can alternatively set API keys, such as ```OPENAI_API_KEY``` in a ```.env``` file and load them.\n",
    "\n",
    "[Note] This is not necessary if you've already set the required API keys in previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load API keys from .env file\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the required libraries\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.output_parsers import PydanticOutputParser\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(temperature=0, model_name=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is an example of an email conversation stored in the variable ```email_conversation``` .\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "email_conversation = \"\"\"\n",
    "From: John (John@bikecorporation.me)\n",
    "To: Kim (Kim@teddyinternational.me)\n",
    "Subject: “ZENESIS” bike distribution cooperation and meeting schedule proposal\n",
    "Dear Mr. Kim,\n",
    "\n",
    "I am John, Senior Executive Director at Bike Corporation. I recently learned about your new bicycle model, \"ZENESIS,\" through your press release. Bike Corporation is a company that leads innovation and quality in the field of bicycle manufacturing and distribution, with long-time experience and expertise in this field.\n",
    "\n",
    "We would like to request a detailed brochure for the ZENESIS model. In particular, we need information on technical specifications, battery performance, and design aspects. This information will help us further refine our proposed distribution strategy and marketing plan.\n",
    "\n",
    "Additionally, to discuss the possibilities for collaboration in more detail, I propose a meeting next Tuesday, January 15th, at 10:00 AM. Would it be possible to meet at your office to have this discussion?\n",
    "\n",
    "Thank you.\n",
    "\n",
    "Best regards,\n",
    "John\n",
    "Senior Executive Director\n",
    "Bike Corporation\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a chain that utilizes the basic ```StrOutputParser```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**Important Parts of the Email:**\n",
      "\n",
      "- **Sender:** John (Senior Executive Director, Bike Corporation)\n",
      "- **Recipient:** Kim (Teddy International)\n",
      "- **Subject:** ZENESIS bike distribution cooperation and meeting schedule proposal\n",
      "- **Request:** Detailed brochure for the ZENESIS model, specifically information on:\n",
      "  - Technical specifications\n",
      "  - Battery performance\n",
      "  - Design aspects\n",
      "- **Purpose:** To refine distribution strategy and marketing plan for ZENESIS.\n",
      "- **Proposed Meeting:** \n",
      "  - Date: Tuesday, January 15th\n",
      "  - Time: 10:00 AM\n",
      "  - Location: Kim's office\n",
      "- **Closing:** Thank you and best regards."
     ]
    }
   ],
   "source": [
    "from itertools import chain\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.messages import AIMessageChunk\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "prompt = PromptTemplate.from_template(\n",
    "    \"Please extract the important parts of the following email.\\n\\n{email_conversation}\"\n",
    ")\n",
    "\n",
    "llm = ChatOpenAI(temperature=0, model_name=\"gpt-4o-mini\")\n",
    "\n",
    "chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "answer = chain.stream({\"email_conversation\": email_conversation})\n",
    "\n",
    "\n",
    "#  A function for real-time output (streaming)\n",
    "def stream_response(response, return_output=False):\n",
    "    \"\"\"\n",
    "    Streams the response from the AI model, processing and printing each chunk.\n",
    "\n",
    "    This function iterates over each item in the 'response' iterable. If an item is an instance of AIMessageChunk, it extracts and prints the content.\n",
    "    If the item is a string, it prints the string directly.\n",
    "    Optionally, the function can return the concatenated string of all response chunks.\n",
    "\n",
    "    Args:\n",
    "    - response (iterable): An iterable of response chunks, which can be AIMessageChunk objects or strings.\n",
    "    - return_output (bool, optional): If True, the function returns the concatenated response string. The default is False.\n",
    "\n",
    "    Returns:\n",
    "    - str: If ```return_output``` is True, the concatenated response string. Otherwise, nothing is returned.\n",
    "    \"\"\"\n",
    "    answer = \"\"\n",
    "    for token in response:\n",
    "        if isinstance(token, AIMessageChunk):\n",
    "            answer += token.content\n",
    "            print(token.content, end=\"\", flush=True)\n",
    "        elif isinstance(token, str):\n",
    "            answer += token\n",
    "            print(token, end=\"\", flush=True)\n",
    "    if return_output:\n",
    "        return answer\n",
    "\n",
    "\n",
    "output = stream_response(answer, return_output=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**Important Parts of the Email:**\n",
      "\n",
      "- **Sender:** John (Senior Executive Director, Bike Corporation)\n",
      "- **Recipient:** Kim (Teddy International)\n",
      "- **Subject:** ZENESIS bike distribution cooperation and meeting schedule proposal\n",
      "- **Request:** Detailed brochure for the ZENESIS model, specifically information on:\n",
      "  - Technical specifications\n",
      "  - Battery performance\n",
      "  - Design aspects\n",
      "- **Purpose:** To refine distribution strategy and marketing plan for ZENESIS.\n",
      "- **Proposed Meeting:** \n",
      "  - Date: Tuesday, January 15th\n",
      "  - Time: 10:00 AM\n",
      "  - Location: Kim's office\n",
      "- **Closing:** Thank you and best regards.\n"
     ]
    }
   ],
   "source": [
    "answer = chain.invoke({\"email_conversation\": email_conversation})\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**Important Parts of the Email:**\n",
      "\n",
      "- **Sender:** John (Senior Executive Director, Bike Corporation)\n",
      "- **Recipient:** Kim (Teddy International)\n",
      "- **Subject:** ZENESIS bike distribution cooperation and meeting schedule proposal\n",
      "- **Request:** Detailed brochure for the ZENESIS model, specifically information on:\n",
      "  - Technical specifications\n",
      "  - Battery performance\n",
      "  - Design aspects\n",
      "- **Purpose:** To refine distribution strategy and marketing plan for ZENESIS.\n",
      "- **Proposed Meeting:** \n",
      "  - Date: Tuesday, January 15th\n",
      "  - Time: 10:00 AM\n",
      "  - Location: Kim's office\n",
      "- **Closing:** Thank you and best regards.\n"
     ]
    }
   ],
   "source": [
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the ```PydanticOutputParser```\n",
    "When provided with email content like the one above, we can parse the email information using the class defined in the ```Pydantic``` style below.\n",
    "\n",
    "For reference, the ```description``` inside the ```Field``` serves as guidance for extracting key information from text-based responses. LLMs rely on this description to extract the required information. Therefore, it is crucial that this description is accurate and clear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmailSummary(BaseModel):\n",
    "    person: str = Field(description=\"The sender of the email\")\n",
    "    email: str = Field(description=\"The email address of the sender\")\n",
    "    subject: str = Field(description=\"The subject of the email\")\n",
    "    summary: str = Field(description=\"A summary of the email content\")\n",
    "    date: str = Field(\n",
    "        description=\"The meeting date and time mentioned in the email content\"\n",
    "    )\n",
    "\n",
    "\n",
    "# Create an output parser using the PydanticOutputParser\n",
    "parser = PydanticOutputParser(pydantic_object=EmailSummary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The output should be formatted as a JSON instance that conforms to the JSON schema below.\n",
      "\n",
      "As an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}\n",
      "the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.\n",
      "\n",
      "Here is the output schema:\n",
      "```\n",
      "{\"properties\": {\"person\": {\"description\": \"The sender of the email\", \"title\": \"Person\", \"type\": \"string\"}, \"email\": {\"description\": \"The email address of the sender\", \"title\": \"Email\", \"type\": \"string\"}, \"subject\": {\"description\": \"The subject of the email\", \"title\": \"Subject\", \"type\": \"string\"}, \"summary\": {\"description\": \"A summary of the email content\", \"title\": \"Summary\", \"type\": \"string\"}, \"date\": {\"description\": \"The meeting date and time mentioned in the email content\", \"title\": \"Date\", \"type\": \"string\"}}, \"required\": [\"person\", \"email\", \"subject\", \"summary\", \"date\"]}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "# Print the instructions.\n",
    "print(parser.get_format_instructions())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a prompt including the following variables:\n",
    "\n",
    "1. ```question```: Receives the user's question.\n",
    "2. ```email_conversation```: Inputs the content of the email conversation.\n",
    "3. ```format```: Specifies the format.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "You are a helpful assistant. \n",
    "\n",
    "QUESTION:\n",
    "{question}\n",
    "\n",
    "EMAIL CONVERSATION:\n",
    "{email_conversation}\n",
    "\n",
    "FORMAT:\n",
    "{format}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "\n",
    "# Apply the partial method and fill in the ```format``` variable in advance\n",
    "prompt = prompt.partial(format=parser.get_format_instructions())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, create a chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = prompt | llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute the chain and review the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```json\n",
      "{\n",
      "  \"person\": \"John\",\n",
      "  \"email\": \"John@bikecorporation.me\",\n",
      "  \"subject\": \"ZENESIS bike distribution cooperation and meeting schedule proposal\",\n",
      "  \"summary\": \"John from Bike Corporation requests a detailed brochure for the ZENESIS bike model, including technical specifications, battery performance, and design aspects. He also proposes a meeting on January 15th at 10:00 AM to discuss collaboration possibilities.\",\n",
      "  \"date\": \"January 15th, 10:00 AM\"\n",
      "}\n",
      "```"
     ]
    }
   ],
   "source": [
    "# Execute the chain and print the result.\n",
    "response = chain.stream(\n",
    "    {\n",
    "        \"email_conversation\": email_conversation,\n",
    "        \"question\": \"Extract the main content of the email.\",\n",
    "    }\n",
    ")\n",
    "\n",
    "# The result is provided in JSON format.\n",
    "output = stream_response(response, return_output=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, use the parser to parse the results and convert them into an ```EmailSummary``` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "person='John' email='John@bikecorporation.me' subject='ZENESIS bike distribution cooperation and meeting schedule proposal' summary='John from Bike Corporation requests a detailed brochure for the ZENESIS bike model, including technical specifications, battery performance, and design aspects. He also proposes a meeting on January 15th at 10:00 AM to discuss collaboration possibilities.' date='January 15th, 10:00 AM'\n"
     ]
    }
   ],
   "source": [
    "structured_output = parser.parse(output)\n",
    "print(structured_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a Chain Including the ```PydanticOutputParser```\n",
    "\n",
    "You can generate the output as a Pydantic object that you define."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reconstruct the entire chain by adding an output parser\n",
    "chain = prompt | llm | parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "person='John' email='John@bikecorporation.me' subject='ZENESIS bike distribution cooperation and meeting schedule proposal' summary='John from Bike Corporation requests a detailed brochure for the ZENESIS bike model, including technical specifications, battery performance, and design aspects. He also proposes a meeting on January 15th at 10:00 AM to discuss collaboration possibilities.' date='January 15th, 10:00 AM'\n"
     ]
    }
   ],
   "source": [
    "# Execute the chain and print the results\n",
    "response = chain.invoke(\n",
    "    {\n",
    "        \"email_conversation\": email_conversation,\n",
    "        \"question\": \"Extract the main content of the email.\",\n",
    "    }\n",
    ")\n",
    "\n",
    "# The results are output in the form of an EmailSummary object\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ```with_structured_output()```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By adding an output parser using ```.with_structured_output(Pydantic)```, you can convert the output into a Pydantic object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_with_structured = ChatOpenAI(\n",
    "    temperature=0, model_name=\"gpt-4o\"\n",
    ").with_structured_output(EmailSummary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EmailSummary(person='John', email='John@bikecorporation.me', subject='“ZENESIS” bike distribution cooperation and meeting schedule proposal', summary='John, Senior Executive Director at Bike Corporation, is interested in the ZENESIS bicycle model and requests a detailed brochure including technical specifications, battery performance, and design aspects. He proposes a meeting to discuss potential collaboration.', date='Tuesday, January 15th, at 10:00 AM')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Call the ```invoke()``` method and print the result\n",
    "answer = llm_with_structured.invoke(email_conversation)\n",
    "answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**\n",
    "\n",
    "One thing to note is that the ```.with_structured_output()``` method does not support the ```stream()``` method."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}